#include "cublas_v2.h"

#include <vector>
#include <cstdio>
void show_matrix(const std::vector<float>& data,const int rows,const int cols){
    for(int r = 0;r<rows;r++){
        for(int c = 0;c<cols;c++){
            printf("data[%d][%d] = %f\n",r,c);
        }
    }

}
int main(int argc, char* argv[])
{
    cublasHandle_t pCuBlas = nullptr;
    float *pCudaB = nullptr;
    float *pCudaA = nullptr;
    float *pCudaC = nullptr;

    cudaError_t cudaStat = cudaSetDevice(0);
    const int m = 2;
    const int k =3;
    const int n = 2;
    cudaStat = cudaMalloc((void**)&pCudaA, m * k * sizeof(float));
    cudaStat = cudaMalloc((void**)&pCudaB, k * n * sizeof(float));
    cudaStat = cudaMalloc((void**)&pCudaC, m * n * sizeof(float));
    cublasStatus_t stat = cublasCreate(&pCuBlas);

    std::vector<float> vA{1,2,3,4,5,6};
    std::vector<float> vB{1,2,3,4,5,6};
    std::vector<float> vC1(4,0);
    std::vector<float> vC2(4,0);
    std::vector<float> vC3(4,0);
    std::vector<float> vC4(4,0);

    cudaMemcpy(pCudaA, vA.data(), sizeof(float)*m*k, cudaMemcpyHostToDevice);
    cudaMemcpy(pCudaB, vB.data(), sizeof(float)*k*n, cudaMemcpyHostToDevice);
    float fAlpha = 1, fBeta = 0;

    // A*B
    stat = cublasSgemm(pCuBlas, 
         CUBLAS_OP_N, CUBLAS_OP_N, 
         2, 2, 3, &fAlpha, pCudaA, 
         m, pCudaB,
         k, &fBeta, pCudaC,
         m
    );
    cudaMemcpy(vC1.data(), pCudaC, sizeof(float)*4, cudaMemcpyDeviceToHost);

    // A 行优先，B列优先
    stat = cublasSgemm(pCuBlas, 
         CUBLAS_OP_T, CUBLAS_OP_N, 
         m, n, k, &fAlpha, pCudaA, 
         n, pCudaB,
         k, &fBeta, pCudaC,
         n
    );

    // A 列优先 B行优先
    cudaMemcpy(vC2.data(), pCudaC, sizeof(float)*4, cudaMemcpyDeviceToHost);
    stat = cublasSgemm(pCuBlas, 
         CUBLAS_OP_N, CUBLAS_OP_T, 
         m, n, k, &fAlpha, pCudaA, 
         2, pCudaB,
         2, &fBeta, pCudaC,
         2
    );
    cudaMemcpy(vC3.data(), pCudaC, sizeof(float)*4, cudaMemcpyDeviceToHost);
    stat = cublasSgemm(pCuBlas, 
         CUBLAS_OP_T, CUBLAS_OP_T, 
         m, n, k, &fAlpha, pCudaA, 
         3, pCudaB,
         2, &fBeta, pCudaC,
         2
    );
    cudaMemcpy(vC4.data(), pCudaC, sizeof(float)*4, cudaMemcpyDeviceToHost);


    for(int i = 0 ;i < vA.size(); i++)
        printf("A %d: %f\n", i, vA[i]);
    for(int i = 0 ;i < vB.size(); i++)
        printf("B %d: %f\n", i, vB[i]);
    for(int i = 0 ;i < vC1.size(); i++)
        printf("C1 %d: %f\n", i, vC1[i]);
    for(int i = 0 ;i < vC2.size(); i++)
        printf("C2 %d: %f\n", i, vC2[i]);
    for(int i = 0 ;i < vC3.size(); i++)
        printf("C3 %d: %f\n", i, vC3[i]);
    for(int i = 0 ;i < vC4.size(); i++)
        printf("C4 %d: %f\n", i, vC4[i]);
        
    if(pCudaA) cudaFree(pCudaA);
    if(pCudaB) cudaFree(pCudaB);
    if(pCudaC) cudaFree(pCudaC);
    if(pCuBlas) cublasDestroy(pCuBlas);

    return 0;

}
